// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/functions/prestosql/width_bucket_array.h>
#include <pollux/expression/expr.h>
#include <pollux/expression/vector_function.h>
#include <pollux/functions/lib/rows_translation_util.h>
#include <pollux/vector/decoded_vector.h>

namespace kumo::pollux::functions {
    namespace {
        template<typename T>
        int64_t widthBucket(
            double operand,
            DecodedVector &elementsHolder,
            int offset,
            int binCount) {
            POLLUX_USER_CHECK_GT(binCount, 0, "Bins cannot be an empty array");
            POLLUX_USER_CHECK(!std::isnan(operand), "Operand cannot be NaN");

            int lower = 0;
            int upper = binCount;
            while (lower < upper) {
                const int index = (lower + upper) / 2;
                POLLUX_USER_CHECK(
                    !elementsHolder.is_null_at(offset + lower) &&
                    !elementsHolder.is_null_at(offset + index) &&
                    !elementsHolder.is_null_at(offset + upper - 1),
                    "Bin values cannot be NULL");

                const auto bin = elementsHolder.value_at<T>(offset + index);
                const auto lowerBin = elementsHolder.value_at<T>(offset + lower);
                const auto upperBin = elementsHolder.value_at<T>(offset + upper - 1);
                POLLUX_USER_CHECK(
                    lowerBin <= bin && bin <= upperBin,
                    "Bin values are not sorted in ascending order");
                POLLUX_USER_CHECK(
                    std::isfinite(bin) && std::isfinite(lowerBin) &&
                    std::isfinite(upperBin),
                    "Bin values must be finite");

                if (operand < bin) {
                    upper = index;
                } else {
                    lower = index + 1;
                }
            }
            return lower;
        }

        template<typename T>
        class WidthBucketArrayFunction : public exec::VectorFunction {
        public:
            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr & /* outputType */,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                context.ensure_writable(rows, BIGINT(), result);
                auto flatResult = result->as_flat_vector<int64_t>()->mutableRawValues();

                exec::DecodedArgs decodedArgs(rows, args, context);
                auto operand = decodedArgs.at(0);
                auto bins = decodedArgs.at(1);

                auto binsArray = bins->base()->as<ArrayVector>();
                auto rawSizes = binsArray->rawSizes();
                auto rawOffsets = binsArray->rawOffsets();
                auto elementsVector = binsArray->elements();
                auto elementsRows = toElementRows(
                    elementsVector->size(),
                    rows,
                    binsArray,
                    bins->nulls(&rows),
                    bins->indices());
                exec::LocalDecodedVector elementsHolder(
                    context, *elementsVector, elementsRows);

                auto indices = bins->indices();
                context.applyToSelectedNoThrow(rows, [&](auto row) {
                    auto size = rawSizes[indices[row]];
                    auto offset = rawOffsets[indices[row]];
                    flatResult[row] = widthBucket<T>(
                        operand->value_at<double>(row), *elementsHolder.get(), offset, size);
                });
            }
        };

        class WidthBucketArrayFunctionConstantBins : public exec::VectorFunction {
        public:
            explicit WidthBucketArrayFunctionConstantBins(std::vector<double> bins)
                : bins_(std::move(bins)) {
            }

            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr & /* outputType */,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                context.ensure_writable(rows, BIGINT(), result);
                auto flatResult = result->as_flat_vector<int64_t>()->mutableRawValues();

                exec::DecodedArgs decodedArgs(rows, args, context);
                auto operand = decodedArgs.at(0);

                context.applyToSelectedNoThrow(rows, [&](auto row) {
                    flatResult[row] = widthBucket(operand->value_at<double>(row), bins_);
                });
            }

        private:
            const std::vector<double> bins_;

            static int64_t widthBucket(double operand, const std::vector<double> &bins) {
                POLLUX_USER_CHECK(!std::isnan(operand), "Operand cannot be NaN");

                int lower = 0;
                int upper = static_cast<int>(bins.size());
                while (lower < upper) {
                    int index = (lower + upper) / 2;
                    auto bin = bins.at(index);

                    if (operand < bin) {
                        upper = index;
                    } else {
                        lower = index + 1;
                    }
                }
                return lower;
            }
        };
    } // namespace

    std::vector<std::shared_ptr<exec::FunctionSignature> >
    widthBucketArraySignature() {
        // double, array(double|bigint) -> bigint
        return {
            exec::FunctionSignatureBuilder()
            .returnType("bigint")
            .argumentType("double")
            .argumentType("array(double)")
            .build(),
            exec::FunctionSignatureBuilder()
            .returnType("bigint")
            .argumentType("double")
            .argumentType("array(bigint)")
            .build(),
        };
    }

    template<typename T>
    std::vector<double> toBinValues(
        const VectorPtr &binsVector,
        vector_size_t offset,
        vector_size_t size) {
        std::vector<double> binValues;
        binValues.reserve(size);
        auto simpleVector = binsVector->as_unchecked<SimpleVector<T> >();

        for (int i = 0; i < size; i++) {
            POLLUX_USER_CHECK(
                !simpleVector->is_null_at(offset + i), "Bin values cannot be null");
            auto value = simpleVector->value_at(offset + i);
            POLLUX_USER_CHECK(std::isfinite(value), "Bin values must be finite");
            if (i > 0) {
                POLLUX_USER_CHECK_GT(
                    value,
                    simpleVector->value_at(offset + i - 1),
                    "Bin values are not sorted in ascending order");
            }
            binValues.push_back(value);
        }

        return binValues;
    }

    std::shared_ptr<exec::VectorFunction> makeWidthBucketArray(
        const std::string &name,
        const std::vector<exec::VectorFunctionArg> &inputArgs,
        const core::QueryConfig & /*config*/) {
        POLLUX_CHECK_EQ(inputArgs.size(), 2);
        const auto &operandVector = inputArgs[0];
        const auto &binsVector = inputArgs[1];

        POLLUX_CHECK_EQ(operandVector.type->kind(), TypeKind::DOUBLE);
        POLLUX_CHECK_EQ(binsVector.type->kind(), TypeKind::ARRAY);

        auto binsTypeKind = binsVector.type->as_array().elementType()->kind();

        auto constantBins = binsVector.constantValue.get();

        try {
            if (constantBins != nullptr && !constantBins->is_null_at(0)) {
                auto binsArrayVector = constantBins->wrapped_vector()->as<ArrayVector>();
                auto binsArrayIndex = constantBins->wrapped_index(0);
                auto size = binsArrayVector->sizeAt(binsArrayIndex);
                POLLUX_USER_CHECK_GT(size, 0, "Bins cannot be an empty array");
                auto offset = binsArrayVector->offsetAt(binsArrayIndex);

                // This is a different behavior comparing to non-constant implementation:
                //
                // In non-constant bins implementation, we only do these checks during
                // binary search, which means we might ignore even if there are infinite
                // value or non-ascending order.
                //
                // In constant bins implementation, we first check bins, so if the bins is
                // invalid, they will fail directly.

                std::vector<double> binValues;
                if (binsTypeKind == TypeKind::DOUBLE) {
                    binValues =
                            toBinValues<double>(binsArrayVector->elements(), offset, size);
                } else if (binsTypeKind == TypeKind::BIGINT) {
                    binValues =
                            toBinValues<int64_t>(binsArrayVector->elements(), offset, size);
                } else {
                    POLLUX_UNSUPPORTED(
                        "Unsupported type of 'bins' argument: {}",
                        binsArrayVector->type()->toString());
                }

                return std::make_shared<WidthBucketArrayFunctionConstantBins>(
                    std::move(binValues));
            }
        } catch (const PolluxException &e) {
            // makeWidthBucketArray should not throw when inputs are invalid.
            VKLOG(1) << e.what();
        }

        if (binsTypeKind == TypeKind::DOUBLE) {
            return std::make_shared<WidthBucketArrayFunction<double> >();
        } else if (binsTypeKind == TypeKind::BIGINT) {
            return std::make_shared<WidthBucketArrayFunction<int64_t> >();
        } else {
            POLLUX_UNSUPPORTED(
                "Unsupported type of 'bins' argument: {}", binsVector.type->toString());
        }
    }
} // namespace kumo::pollux::functions
